package animation;

import com.crunch.math.MathOps;
import com.crunch.math.Matrix33af;
import com.crunch.math.Vector2f;

import java.util.ArrayList;
import java.util.BitSet;

/**
 * Used to evaluate animations.
 */
public class AnimationEvaluator {
	public AnimationEvaluator() {
		transformations = new ArrayList<Transformation>();
		activeBones = new BitSet();
	}

	// computes the bone matrices for the character state
	public void evaluate(Character.State state) {
		Skeleton skeleton = state.getCharacter().getSkeleton();

		// allocate enough space for transformations and active bone bitfield
		int boneCount = skeleton.getBoneCount();
		int requiredDepth = Math.max(1, state.getCharacter().getAnimationTree().getEvaluatorDepthRequirement());
		int requiredSize = boneCount * requiredDepth;
		int currentSize = transformations.size();
		if (currentSize < requiredSize) {
			transformations.ensureCapacity(requiredSize);
			for (int i = currentSize; i < requiredSize; ++i) {
				transformations.add(new Transformation());
			}
			// make sure bitset is large enough by setting the last bit (then clearing)
			activeBones.set(requiredSize - 1);
			activeBones.clear(requiredSize - 1);
		}

		// initialize to the default pose
		for (int i = 0; i < boneCount; ++i) {
			Transformation tf = transformations.get(i);
			Skeleton.Bone bone = state.getCharacter().getSkeleton().getBone(i);
			tf.translation.set(bone.getTranslation());
			tf.rotation = bone.getRotation();
			tf.scale.set(bone.getScale());
			// point last active transformation to self
			tf.lastActiveTransformation = i;
		}

		// evaluate the animation tree
		if (state.getCharacter().getAnimationTree().getAnimationTree().getRoot() != null) {
			evaluate(state, 0, 0, 0);
		}

		// recursively multiply poses to get final result
		for (int i = 0; i < skeleton.getRootBoneCount(); ++i) {
			computePoseMatrices(skeleton, skeleton.getRootBone(i), state.getBoneMatrices());
		}
	}

	private void evaluate(Character.State state, int nodeIndex, int bufferOffset, int currentPoseOffset) {
		// clear this node's bone indices
		activeBones.clear(bufferOffset, bufferOffset + state.getCharacter().getSkeleton().getBoneCount());
		// prepare the buffer by setting the last active pose pointers
		for (int i = 0; i < state.getCharacter().getSkeleton().getBoneCount(); ++i) {
			transformations.get(bufferOffset + i).lastActiveTransformation =
					transformations.get(currentPoseOffset + i).lastActiveTransformation;
		}

		AnimationTreeInstance.Node node = state.getCharacter().getAnimationTree().getNode(nodeIndex);
		switch (node.getNode().getType()) {
			case ANIMATION:
				evaluateAnimationNode(state, nodeIndex, bufferOffset);
				break;
			case SELECT:
				evaluateSelectNode(state, nodeIndex, bufferOffset, currentPoseOffset);
				break;
			case BLEND:
				evaluateBlendNode(state, nodeIndex, bufferOffset, currentPoseOffset);
				break;
			case SEQUENCE:
				evaluateSequenceNode(state, nodeIndex, bufferOffset);
				break;
			default:
				throw new IllegalArgumentException("Unknown node type " + node.getNode().getType());
		}
	}

	private static float getKeyframeRatio(Animation.AnimatedBone bone, int kfA, int kfB, float time) {
		return (kfA == kfB) ? 0.0f :
				(time - bone.getKeyframe(kfA).getTime()) / (bone.getKeyframe(kfB).getTime() - bone.getKeyframe(kfA).getTime());
	}

	private void evaluateAnimationNode(Character.State state, int nodeIndex, int bufferOffset) {
		AnimationTreeInstance.AnimationNode nodeInst =
				(AnimationTreeInstance.AnimationNode) state.getCharacter().getAnimationTree().getNode(nodeIndex);
		AnimationTreeInstance.State.AnimationNodeState nodeState =
				(AnimationTreeInstance.State.AnimationNodeState) state.getAnimationTreeState().getNodeState(nodeIndex);

		AnimationInstance animationInst = nodeInst.getAnimation();
		Animation animation = animationInst.getAnimation();
		// set active bones for each bone included in the animation
		for (int i = 0; i < animation.getAnimatedBoneCount(); ++i) {
			activeBones.set(bufferOffset + animationInst.getAnimatedBoneIndex(i));
		}

		float time = nodeState.getTime();

		// find the fade ratio
		float fadeInRatio = (animation.getFadeInTime() == 0.0f) ? 1.0f :
				Math.max(0.0f, Math.min(1.0f, time / animation.getFadeInTime()));
		float fadeOutRatio = (animation.getFadeOutTime() == 0.0f) ? 1.0f :
				Math.max(0.0f, Math.min(1.0f, (animation.getLength() - time) / animation.getFadeOutTime()));
		float fadeRatio = Math.min(fadeInRatio, fadeOutRatio);

		// compute each bone
		if (fadeRatio == 0.0f) {
			// take directly from previous level
			for (int i = 0; i < animation.getAnimatedBoneCount(); ++i) {
				int boneId = animationInst.getAnimatedBoneIndex(i);

				Transformation newPose = transformations.get(bufferOffset + boneId);
				Transformation activePose = transformations.get(newPose.lastActiveTransformation);

				newPose.translation.set(activePose.translation);
				newPose.rotation = activePose.rotation;
				newPose.scale = activePose.scale;
			}
		} else if (fadeRatio == 1.0f) {
			// apply animation fully
			for (int i = 0; i < animation.getAnimatedBoneCount(); ++i) {
				int boneId = animationInst.getAnimatedBoneIndex(i);
				Animation.AnimatedBone bone = animation.getAnimatedBone(i);

				if (bone.getKeyframeCount() > 0) {
					Transformation newPose = transformations.get(bufferOffset + boneId);

					int kfA = bone.getLowerKeyframe(time);
					int kfB = kfA + 1;
					kfA = Math.max(kfA, 0);
					kfB = Math.min(kfB, bone.getKeyframeCount()-1);
					float ratio = getKeyframeRatio(bone, kfA, kfB, time);
					Animation.Keyframe keyframeA = bone.getKeyframe(kfA);
					Animation.Keyframe keyframeB = bone.getKeyframe(kfB);

					newPose.translation.set(MathOps.lerp(keyframeA.getTranslation(), keyframeB.getTranslation(), ratio));
					newPose.rotation = MathOps.lerpAngles(keyframeA.getRotation(), keyframeB.getRotation(), ratio);
					newPose.scale.set(MathOps.lerp(keyframeA.getScale(), keyframeB.getScale(), ratio));
				}
			}
		} else {
			// apply animation and blend with previous level
			for (int i = 0; i < animation.getAnimatedBoneCount(); ++i) {
				int boneId = animationInst.getAnimatedBoneIndex(i);
				Animation.AnimatedBone bone = animation.getAnimatedBone(i);

				if (bone.getKeyframeCount() > 0) {
					Transformation newPose = transformations.get(bufferOffset + boneId);
					Transformation activePose = transformations.get(newPose.lastActiveTransformation);

					int kfA = bone.getLowerKeyframe(time);
					int kfB = kfA + 1;
					kfA = Math.max(kfA, 0);
					kfB = Math.min(kfB, bone.getKeyframeCount()-1);
					float ratio = getKeyframeRatio(bone, kfA, kfB, time);
					Animation.Keyframe keyframeA = bone.getKeyframe(kfA);
					Animation.Keyframe keyframeB = bone.getKeyframe(kfB);

					Vector2f translation = MathOps.lerp(keyframeA.getTranslation(), keyframeB.getTranslation(), ratio);
					float rotation = MathOps.lerpAngles(keyframeA.getRotation(), keyframeB.getRotation(), ratio);
					Vector2f scale = MathOps.lerp(keyframeA.getScale(), keyframeB.getScale(), ratio);

					newPose.translation.set(MathOps.lerp(activePose.translation, translation, fadeRatio));
					newPose.rotation = MathOps.lerpAngles(activePose.rotation, rotation, fadeRatio);
					newPose.scale.set(MathOps.lerp(activePose.scale, scale, fadeRatio));
				}
			}
		}
	}

	private void evaluateSelectNode(Character.State state, int nodeIndex, int bufferOffset, int currentPoseOffset) {
		AnimationTreeInstance.SelectNode nodeInst =
				(AnimationTreeInstance.SelectNode) state.getCharacter().getAnimationTree().getNode(nodeIndex);
		AnimationTree.SelectNode node = (AnimationTree.SelectNode) nodeInst.getNode();
		AnimationTreeInstance.State.SelectNodeState nodeState =
				(AnimationTreeInstance.State.SelectNodeState) state.getAnimationTreeState().getNodeState(nodeIndex);

		if (nodeState.getIndex() >= 0 && nodeState.getIndex() < node.getChildCount()) {
			// simply "pass through" to the selected child node
			evaluate(state, nodeInst.getChildNodeIndex(nodeState.getIndex()), bufferOffset, currentPoseOffset);
		}
		// otherwise, invalid index - don't apply animation at all
	}

	private void evaluateBlendNode(Character.State state, int nodeIndex, int bufferOffset, int currentPoseOffset) {
		AnimationTreeInstance.BlendNode nodeInst =
				(AnimationTreeInstance.BlendNode) state.getCharacter().getAnimationTree().getNode(nodeIndex);
		AnimationTree.BlendNode node = (AnimationTree.BlendNode) nodeInst.getNode();
		AnimationTreeInstance.State.BlendNodeState nodeState =
				(AnimationTreeInstance.State.BlendNodeState) state.getAnimationTreeState().getNodeState(nodeIndex);

		if (node.getChildCount() == 0) {
			return; // do nothing
		}

		float weight = nodeState.getWeight();
		if (node.getWrap()) {
			float maxWeight = (float) node.getChildCount();
			weight = Math.max(0.0f, Math.min(maxWeight, weight));
			if (weight == maxWeight) {
				weight = 0.0f;
			}
		} else {
			weight = Math.max(0.0f, Math.min((float) (node.getChildCount()-1), weight));
		}

		// break weight up into integer part (child index) and weight part
		int leftChildIndex = (int) Math.floor(weight);
		int rightChildIndex = (leftChildIndex + 1 == node.getChildCount()) ? 0 : leftChildIndex + 1;
		assert(leftChildIndex < node.getChildCount()); // check for rounding issue
		weight -= (float) leftChildIndex;

		// optimize for cases of 0 weight
		// we will never actually get a weight of 1 (since it's equivalent to a 0 weight case)
		if (weight == 0.0f) {
			// simply pass through, like a select node
			int child = nodeInst.getChildNodeIndex(leftChildIndex);
			evaluate(state, child, bufferOffset, currentPoseOffset);
		} else {
			int leftChild = nodeInst.getChildNodeIndex(leftChildIndex);
			int rightChild = nodeInst.getChildNodeIndex(rightChildIndex);

			int leftBufferOffset = bufferOffset + state.getCharacter().getSkeleton().getBoneCount();
			int rightBufferOffset = leftBufferOffset + state.getCharacter().getSkeleton().getBoneCount();

			// evaluate left and right nodes into buffer slots
			evaluate(state, leftChild, leftBufferOffset, currentPoseOffset);
			evaluate(state, rightChild, rightBufferOffset, currentPoseOffset);

			// blend between the two results
			for (int i = 0; i < state.getCharacter().getSkeleton().getBoneCount(); ++i) {
				boolean leftActive = activeBones.get(leftBufferOffset + i);
				boolean rightActive = activeBones.get(rightBufferOffset + i);

				if (leftActive || rightActive) {
					Transformation newPose = transformations.get(bufferOffset + i);
					Transformation leftPose;
					Transformation rightPose;

					if (!rightActive) {
						// only left is active - treat active pose as right child
						leftPose = transformations.get(leftBufferOffset + i);
						rightPose = transformations.get(newPose.lastActiveTransformation);
					} else if (!leftActive) {
						// only right is active - treat active pose as left child
						leftPose = transformations.get(newPose.lastActiveTransformation);
						rightPose = transformations.get(rightBufferOffset + i);
					} else {
						// both are active - blend between both bones
						leftPose = transformations.get(leftBufferOffset + i);
						rightPose = transformations.get(rightBufferOffset + i);
					}

					// blend between child poses
					newPose.translation.set(MathOps.lerp(leftPose.translation, rightPose.translation, weight));
					newPose.rotation = MathOps.lerpAngles(leftPose.rotation, rightPose.rotation, weight);
					newPose.scale.set(MathOps.lerp(leftPose.scale, rightPose.scale, weight));
				}

				// bone is active if either side activated it
				activeBones.set(bufferOffset + i, leftActive || rightActive);
			}
		}
	}

	private void evaluateSequenceNode(Character.State state, int nodeIndex, int bufferOffset) {
		AnimationTreeInstance.SequenceNode nodeInst =
				(AnimationTreeInstance.SequenceNode) state.getCharacter().getAnimationTree().getNode(nodeIndex);
		AnimationTree.SequenceNode node = (AnimationTree.SequenceNode) nodeInst.getNode();
		AnimationTreeInstance.State.SequenceNodeState nodeState =
				(AnimationTreeInstance.State.SequenceNodeState) state.getAnimationTreeState().getNodeState(nodeIndex);

		int boneCount = state.getCharacter().getSkeleton().getBoneCount();

		// evaluate the child nodes in order
		for (int i = 0; i < node.getChildCount(); ++i) {
			int child = nodeInst.getChildNodeIndex(i);
			int childBufferOffset = bufferOffset + state.getCharacter().getSkeleton().getBoneCount();

			// pass in the this node's buffer offset as the current pose
			// this is because as nodes are evaluated in sequence, the current pose (at this sequence node) is modified
			// later children should blend with "later" current poses

			float weight = nodeState.getWeight(i);
			// don't need to do anything if weight is 0
			if (weight != 0.0f) {
				// evaluate the child node
				evaluate(state, child, childBufferOffset, bufferOffset);

				if (weight == 1.0f) {
					// optimization - don't blend, just set active bones directly
					for (int b = 0; b < boneCount; ++b) {
						boolean active = activeBones.get(childBufferOffset + b);
						if (active) {
							// copy transformation
							Transformation tf = transformations.get(bufferOffset + b);
							Transformation childTf = transformations.get(childBufferOffset + b);

							tf.translation.set(childTf.translation);
							tf.rotation = childTf.rotation;
							tf.scale.set(childTf.scale);
							tf.lastActiveTransformation = bufferOffset + b;
							// this bone is now active
							activeBones.set(bufferOffset + b);
						}
					}
				} else {
					// blend between current pose and child pose
					for (int b = 0; b < boneCount; ++b) {
						boolean active = activeBones.get(childBufferOffset + b);
						if (active) {
							// blend with last active transformation
							Transformation tf = transformations.get(bufferOffset + b);
							Transformation latf = transformations.get(tf.lastActiveTransformation);
							Transformation childTf = transformations.get(childBufferOffset + b);

							tf.translation.set(MathOps.lerp(latf.translation, childTf.translation, weight));
							tf.rotation = MathOps.lerpAngles(latf.rotation, childTf.rotation, weight);
							tf.scale.set(MathOps.lerp(latf.scale, childTf.scale, weight));
							tf.lastActiveTransformation = bufferOffset + b;
							// this bone is now active
							activeBones.set(bufferOffset + b);
						}
					}
				}
			}
		}
	}

	private void computePoseMatrices(Skeleton skeleton, int bone, Matrix33af[] poseMatrices) {
		Transformation tf = transformations.get(bone);
		Matrix33af mat = poseMatrices[bone];
		Skeleton.Bone b = skeleton.getBone(bone);

		mat.setIdentity();
		mat.scaleAssign(tf.scale);
		mat.rotateAssign(tf.rotation);
		mat.translateAssign(tf.translation);
		if (b.getParent() >= 0) {
			// multiply with parent's matrix
			mat.preMultiplyAssign(poseMatrices[b.getParent()]);
		}

		for (int i = 0; i < b.getChildCount(); ++i) {
			computePoseMatrices(skeleton, b.getChild(i), poseMatrices);
		}
	}

	private static class Transformation {
		public Vector2f translation = new Vector2f();
		public float rotation;
		public Vector2f scale = new Vector2f(1.0f, 1.0f);
		public int lastActiveTransformation; // points to last active pose
	}
	private ArrayList<Transformation> transformations;
	private BitSet activeBones;
}